/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2020-2022. All rights reserved.
 * Description: 集合通信算子信息库头文件
 * Author: jiangchen
 * Create: 2020-5-15
 */

#ifndef OPS_KERNEL_INFO_STORE__BASEH
#define OPS_KERNEL_INFO_STORE__BASEH

#include "comm.h"
#include "hcom_ops_stores.h"
#include "base.h"
#include "common/opskernel/ge_task_info.h"
#include "common/opskernel/ops_kernel_info_types.h"
#include "graph/node.h"
#include "proto/task.pb.h"
#include "op_common_def.h"

namespace hccl {
// Ge适配的类
class HCCLOpsKernelInfoStore : public ge::OpsKernelInfoStore {
public:
    HCCLOpsKernelInfoStore();
    ~HCCLOpsKernelInfoStore() override;
    // initialize opsKernelInfoStore
    ge::Status Initialize(const map<string, string> &options) override;
    // close opsKernelInfoStore
    ge::Status Finalize() override;
    // get all opsKernelInfo
    void GetAllOpsKernelInfo(map<string, ge::OpInfo> &infos) const override;
    // whether the opsKernelInfoStore is supported based on the operator attribute
    bool CheckSupported(const ge::OpDescPtr &opDescPtr, std::string& unSupportedReason) const override;
protected:
    virtual HcclResult SetCustomKernelInfo(ge::OpInfo &opinfo, std::map<string, ge::OpInfo> &infos) const = 0;
    virtual HcclResult GetSupportedOP(std::vector<std::string> &hcclSupportOp) const = 0;
    virtual HcclResult GetDataTypeFromTaskInfo(const ge::GETaskInfo &task, HcclDataType &dataType) const = 0;

    HcclResult CheckSupportedOP(const std::string &sCollectiveType) const;
    HcclResult GetCollectiveTypeFromTaskInfo(const ge::GETaskKernelHcclInfo &hcclInfo, std::string &sCollectiveType);
    HcclResult GetCountFromTaskInfo(const ge::GETaskKernelHcclInfo &hcclInfo, u64 &count);
    HcclResult GetStreamMainFromTaskInfo(const ge::GETaskInfo &taskDef, rtStream_t &stream);
    HcclResult GetInputAddrFromTaskInfo(const ge::GETaskKernelHcclInfo &hcclInfo, uintptr_t &inputAddr);
    HcclResult GetOutputAddrFromTaskInfo(const ge::GETaskKernelHcclInfo &hcclInfo, uintptr_t &outputAddr);
    HcclResult GetRootFromTaskInfo(const ge::GETaskKernelHcclInfo &hcclInfo, u32 &root);
    HcclResult GetReduceTypeFromTaskInfo(const ge::GETaskKernelHcclInfo &hcclInfo, HcclReduceOp &opType);
    HcclResult GetGlobalWorkSpaceAddrFromTaskInfo(const ge::GETaskKernelHcclInfo &hcclInfo,
        std::vector<void *> &globalWorkSpaceAddr);
};
}
#endif // OPS_KERNEL_INFO_STORE__BASEH
